package corpus;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;

import parser.ParserAnnotation;
import parser.ParserPipeline;
import parser.StanfordPipeline;

public class NGram {
	private BufferedReader in;
	private ParserPipeline pipeline;

	public static void main(String[] args) throws IOException {
		ParserPipeline pipeline = new StanfordPipeline("ssplit, tokenize", "");
		File corpus = new File("qv.2.txt");
		NGram n = new NGram(corpus, pipeline);
		n.makeNGram(new File("bigram.txt"), 2);
	}

	public NGram(File corpus, ParserPipeline pipeline)
			throws FileNotFoundException {
		in = new BufferedReader(new FileReader(corpus));
		this.pipeline = pipeline;
	}

	public void makeNGram(File out, int n) throws IOException {
		HashMap<String, Integer> words = new HashMap<String, Integer>();
		int lineCount = 0;
		int totalWords = 0;
		while (true) {
			String line = in.readLine();
			if (line == null)
				break;
			if (line.isEmpty())
				continue;

			ParserAnnotation pa = pipeline.make(line);
			pa.parse();

			for (int sent = 0; sent < pa.getNbrSentences(); sent++) {
				List<String> tokens = pa.getTokens(sent);
				for (int i = 0; i < tokens.size(); i++) {
					String s = "";
					for (int j = 0; j < n; j++) {
						if (j - n < 0) {
							s += "<BEGIN> ";
						} else {
							s += tokens.get(j - n);
						}
					}
					s += tokens.get(i);
					totalWords++;
					if (words.containsKey(s)) {
						words.put(s, words.get(s) + 1);
					} else {
						words.put(s, 1);
					}
				}
			}

			lineCount++;
			if (lineCount > 1000)
				break;
			if (lineCount % 1000 == 0)
				System.out.println(lineCount);
		}

		BufferedWriter fw = new BufferedWriter(new FileWriter(out));
		for (Entry<String, Integer> e : words.entrySet()) {
			fw.append(e.getKey() + " " + e.getValue() + "\n");
		}
		fw.append("Total words: " + totalWords + "\n");

		fw.flush();
		fw.close();
	}
}
